#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
修复数据库表结构
"""

import sqlite3
import os

def get_db_path():
    """获取数据库路径"""
    current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    db_path = os.path.join(current_dir, "..", "flux_prompt_helper", "promptManager-server", "instance", "prompt_manager.db")
    return db_path

def fix_database():
    """修复数据库表结构"""
    db_path = get_db_path()
    
    if not os.path.exists(os.path.dirname(db_path)):
        os.makedirs(os.path.dirname(db_path), exist_ok=True)
    
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    try:
        # 删除旧的categories表（如果存在）
        cursor.execute("DROP TABLE IF EXISTS categories")
        print("删除旧的categories表")
        
        # 创建新的categories表
        cursor.execute("""
            CREATE TABLE categories (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                category_key TEXT NOT NULL UNIQUE,
                category_name TEXT NOT NULL,
                description TEXT DEFAULT '',
                parent_id INTEGER DEFAULT NULL,
                level INTEGER DEFAULT 0,
                data TEXT NOT NULL,
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (parent_id) REFERENCES categories (id)
            )
        """)
        print("创建新的categories表")
        
        # 创建索引
        cursor.execute("CREATE INDEX idx_categories_key ON categories (category_key)")
        cursor.execute("CREATE INDEX idx_categories_parent ON categories (parent_id)")
        cursor.execute("CREATE INDEX idx_categories_level ON categories (level)")
        print("创建索引")
        
        # 提交事务
        conn.commit()
        print("数据库修复完成!")
        
        return True
        
    except Exception as e:
        print(f"修复失败: {e}")
        conn.rollback()
        return False
    finally:
        conn.close()

if __name__ == "__main__":
    fix_database()
